import numpy as np

def jacobi_poisson_solver(source_term, dx, dy, boundary_conditions, tol=1e-7, max_iters=1000000 , psi_init=None):
    """
    Solve the Poisson equation using Jacobi iteration.
    
    Parameters:
    - source_term: The source term (right-hand side of the Poisson equation), shape (Ny, Nx).
    - dx, dy: Grid spacings in x and y directions.
    - boundary_conditions: A function that defines boundary conditions for psi.
                           It should take 'psi' and 'idx' (0=top, 1=bottom, 2=left, 3=right) and return the BC values.
    - tol: Tolerance for convergence.
    - max_iters: Maximum number of iterations.
    - psi_init: Initial guess for psi, default is zeros.
    
    Returns:
    - psi: The computed solution of the Poisson equation.
    """
    Ny, Nx = source_term.shape
    if psi_init is None:
        psi = np.zeros_like(source_term)
    else:
        psi = psi_init.copy()

    for k in range(max_iters):
        if k%1000==0:
            print(k)

        psi_new = np.zeros_like(psi)
        
        # Update interior points
        psi_new[1:-1, 1:-1] = (-source_term[1:-1, 1:-1] 
                               + (psi[2:, 1:-1] + psi[:-2, 1:-1]) / dy**2 
                               + (psi[1:-1, 2:] + psi[1:-1, :-2]) / dx**2) / (2 / dx**2 + 2 / dy**2)
        
        # Apply boundary conditions
        psi_new = boundary_conditions(psi_new, 0)  # Top
        psi_new = boundary_conditions(psi_new, 1)  # Bottom
        psi_new = boundary_conditions(psi_new, 2)  # Left
        psi_new = boundary_conditions(psi_new, 3)  # Right
        
        # Check for convergence
        if np.max(np.abs(psi_new - psi)) < tol:
            break
            
        psi = psi_new.copy()

    return psi

# 示例边界条件函数
def example_boundary_conditions(psi, idx):
    Ny, Nx = psi.shape
    if idx == 0:  # Top
        psi[-1, :] = psi[-2, :]  # 或其他适当的边界条件
    elif idx == 1:  # Bottom
        psi[0, :] = 0  # 或其他
    elif idx == 2:  # Left
        psi[:, 0] = 0  # 或其他
    elif idx == 3:  # Right
        psi[:, -1] = 0  # 或其他
    return psi

def analytic_solution(x, y):
    """解析解"""
    return np.sin(np.pi * x) * np.sin(np.pi * y)

def source_term(x, y, dx, dy):
    """根据解析解计算源项"""
    nx, ny = x.shape[0], y.shape[1]
    f = np.zeros((nx, ny))
    for i in range(nx):
        for j in range(ny):
            f[i, j] = -2*np.pi**2 * np.sin(np.pi * x[i, j]) * np.sin(np.pi * y[i, j])
    return f

def zero_dirichlet_bc(psi, idx):
    """全零Dirichlet边界条件"""
    Ny, Nx = psi.shape
    if idx == 0:  # Top
        psi[-1, :] = 0
    elif idx == 1:  # Bottom
        psi[0, :] = 0
    elif idx == 2:  # Left
        psi[:, 0] = 0
    elif idx == 3:  # Right
        psi[:, -1] = 0
    return psi

if __name__ == "__main__":
    Nx, Ny = 100, 100  # 网格大小
    Lx, Ly = 2.0, 2.0  # 域的尺寸
    dx, dy = Lx / (Nx - 1), Ly / (Ny - 1)
    x = np.linspace(-Lx/2, Lx/2, Nx)
    y = np.linspace(-Ly/2, Ly/2, Ny)
    X, Y = np.meshgrid(x, y)

    # 计算源项
    source = source_term(X, Y, dx, dy)

    # 使用求解器求解
    psi_numerical = jacobi_poisson_solver(source, dx, dy, zero_dirichlet_bc, tol=1e-11)

    # 计算解析解
    psi_analytic = analytic_solution(X, Y)

    # 计算并打印误差
    error = np.max(np.abs(psi_numerical - psi_analytic))
    print(f"Maximum error between numerical and analytic solution: {error}")

    # 可视化结果（可选）
    import matplotlib.pyplot as plt
    fig, axs = plt.subplots(1, 2, figsize=(12, 6))
    cax1 = axs[0].imshow(psi_numerical, extent=[0, Lx, 0, Ly], cmap='viridis', origin='lower')
    axs[0].set_title('Numerical Solution')
    fig.colorbar(cax1, ax=axs[0])

    cax2 = axs[1].imshow(psi_analytic, extent=[0, Lx, 0, Ly], cmap='viridis', origin='lower')
    axs[1].set_title('Analytic Solution')
    fig.colorbar(cax2, ax=axs[1])

    plt.show()